//! 查询引擎实现
//!
//! 基于DataFusion的高性能查询引擎

use std::sync::Arc;
use std::collections::HashMap;
use datafusion::prelude::*;
use datafusion::catalog::TableProvider;
use datafusion::arrow::record_batch::RecordBatch;
use datafusion::arrow::array::{Array, StringArray, Int64Array, Float64Array};
use datafusion::arrow::datatypes::{Schema, Field, DataType};
use serde_json::Value;
use chrono::Utc;

use data_gateway_kernel::{Query, QueryResult, QueryType};
use data_gateway_kernel::query::{QueryStats, QueryStatus, QueryResultData};

/// 查询引擎
pub struct QueryEngine {
    ctx: SessionContext,
    registered_tables: HashMap<String, Arc<dyn TableProvider>>,
}

impl QueryEngine {
    /// 创建新的查询引擎
    pub fn new() -> Self {
        let ctx = SessionContext::new();
        Self {
            ctx,
            registered_tables: HashMap::new(),
        }
    }

    /// 注册表
    pub async fn register_table(&mut self, name: &str, table: Arc<dyn TableProvider>) -> anyhow::Result<()> {
        self.ctx.register_table(name, table.clone())?;
        self.registered_tables.insert(name.to_string(), table);
        Ok(())
    }

    /// 执行SQL查询
    pub async fn execute_sql(&self, sql: &str) -> anyhow::Result<Vec<RecordBatch>> {
        let df = self.ctx.sql(sql).await?;
        let batches = df.collect().await?;
        Ok(batches)
    }

    /// 执行查询
    pub async fn execute(&self, query: Query) -> anyhow::Result<QueryResult> {
        let start_time = std::time::Instant::now();

        match query.query_type {
            QueryType::Sql => {
                match self.execute_sql(&query.content).await {
                    Ok(batches) => {
                        let duration = start_time.elapsed();
                        let rows_processed: usize = batches.iter().map(|b| b.num_rows()).sum();

                        Ok(QueryResult {
                            query_id: query.id,
                            status: QueryStatus::Success,
                            data: QueryResultData::Arrow(batches),
                            stats: QueryStats {
                                start_time: Utc::now(),
                                end_time: Some(Utc::now()),
                                duration_ms: Some(duration.as_millis() as u64),
                                rows_processed: Some(rows_processed),
                                rows_returned: Some(rows_processed),
                                memory_used: Some(0),
                                cpu_time_ms: Some(duration.as_millis() as u64),
                                cache_hit: false,
                            },
                            error: None,
                        })
                    }
                    Err(e) => {
                        let duration = start_time.elapsed();
                        Ok(QueryResult::failed(
                            query.id,
                            e.to_string(),
                            QueryStats {
                                start_time: Utc::now(),
                                end_time: Some(Utc::now()),
                                duration_ms: Some(duration.as_millis() as u64),
                                rows_processed: Some(0),
                                rows_returned: Some(0),
                                memory_used: Some(0),
                                cpu_time_ms: Some(duration.as_millis() as u64),
                                cache_hit: false,
                            },
                        ))
                    }
                }
            }
            _ => {
                Ok(QueryResult::failed(
                    query.id,
                    "不支持的查询类型".to_string(),
                    QueryStats {
                        start_time: Utc::now(),
                        end_time: Some(Utc::now()),
                        duration_ms: Some(0),
                        rows_processed: Some(0),
                        rows_returned: Some(0),
                        memory_used: Some(0),
                        cpu_time_ms: Some(0),
                        cache_hit: false,
                    },
                ))
            }
        }
    }

    /// 创建内存表
    pub async fn create_memory_table(&mut self, name: &str, data: Vec<HashMap<String, Value>>) -> anyhow::Result<()> {
        if data.is_empty() {
            return Ok(());
        }

        // 推断schema
        let first_row = &data[0];
        let mut fields = Vec::new();

        for (key, value) in first_row {
            let data_type = match value {
                Value::String(_) => DataType::Utf8,
                Value::Number(n) if n.is_i64() => DataType::Int64,
                Value::Number(_) => DataType::Float64,
                Value::Bool(_) => DataType::Boolean,
                _ => DataType::Utf8, // 默认为字符串
            };
            fields.push(Field::new(key, data_type, true));
        }

        let schema = Arc::new(Schema::new(fields));

        // 创建数组
        let mut arrays: Vec<Arc<dyn Array>> = Vec::new();

        for field in schema.fields() {
            let column_data: Vec<_> = data.iter()
                .map(|row| row.get(field.name()).cloned().unwrap_or(Value::Null))
                .collect();

            let array: Arc<dyn Array> = match field.data_type() {
                DataType::Utf8 => {
                    let string_data: Vec<Option<String>> = column_data.iter()
                        .map(|v| match v {
                            Value::String(s) => Some(s.clone()),
                            Value::Null => None,
                            _ => Some(v.to_string()),
                        })
                        .collect();
                    Arc::new(StringArray::from(string_data))
                }
                DataType::Int64 => {
                    let int_data: Vec<Option<i64>> = column_data.iter()
                        .map(|v| match v {
                            Value::Number(n) => n.as_i64(),
                            Value::Null => None,
                            _ => None,
                        })
                        .collect();
                    Arc::new(Int64Array::from(int_data))
                }
                DataType::Float64 => {
                    let float_data: Vec<Option<f64>> = column_data.iter()
                        .map(|v| match v {
                            Value::Number(n) => n.as_f64(),
                            Value::Null => None,
                            _ => None,
                        })
                        .collect();
                    Arc::new(Float64Array::from(float_data))
                }
                _ => {
                    let string_data: Vec<Option<String>> = column_data.iter()
                        .map(|v| match v {
                            Value::Null => None,
                            _ => Some(v.to_string()),
                        })
                        .collect();
                    Arc::new(StringArray::from(string_data))
                }
            };
            arrays.push(array);
        }

        let batch = RecordBatch::try_new(schema, arrays)?;
        let table = datafusion::datasource::memory::MemTable::try_new(
            batch.schema(),
            vec![vec![batch]]
        )?;

        self.register_table(name, Arc::new(table)).await?;
        Ok(())
    }

    /// 获取已注册的表列表
    pub fn list_tables(&self) -> Vec<String> {
        self.registered_tables.keys().cloned().collect()
    }

    /// 获取表的schema
    pub async fn get_table_schema(&self, table_name: &str) -> anyhow::Result<Schema> {
        if let Some(table) = self.registered_tables.get(table_name) {
            Ok(table.schema().as_ref().clone())
        } else {
            Err(anyhow::anyhow!("表不存在: {}", table_name))
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::collections::HashMap;
    use serde_json::json;
    use data_gateway_kernel::query::{QueryOptions, QueryContext};

    #[tokio::test]
    async fn test_query_engine_creation() {
        let engine = QueryEngine::new();
        assert!(engine.list_tables().is_empty());
    }

    #[tokio::test]
    async fn test_create_memory_table() {
        let mut engine = QueryEngine::new();

        let data = vec![
            {
                let mut row = HashMap::new();
                row.insert("id".to_string(), json!(1));
                row.insert("name".to_string(), json!("Alice"));
                row.insert("age".to_string(), json!(30));
                row
            },
            {
                let mut row = HashMap::new();
                row.insert("id".to_string(), json!(2));
                row.insert("name".to_string(), json!("Bob"));
                row.insert("age".to_string(), json!(25));
                row
            },
        ];

        engine.create_memory_table("users", data).await.unwrap();

        let tables = engine.list_tables();
        assert_eq!(tables.len(), 1);
        assert!(tables.contains(&"users".to_string()));
    }

    #[tokio::test]
    async fn test_sql_query_execution() {
        let mut engine = QueryEngine::new();

        // 创建测试数据
        let data = vec![
            {
                let mut row = HashMap::new();
                row.insert("id".to_string(), json!(1));
                row.insert("name".to_string(), json!("Alice"));
                row.insert("score".to_string(), json!(95.5));
                row
            },
            {
                let mut row = HashMap::new();
                row.insert("id".to_string(), json!(2));
                row.insert("name".to_string(), json!("Bob"));
                row.insert("score".to_string(), json!(87.2));
                row
            },
        ];

        engine.create_memory_table("students", data).await.unwrap();

        // 执行SQL查询
        let query = Query {
            id: "test_query".to_string(),
            query_type: QueryType::Sql,
            content: "SELECT name, score FROM students WHERE score > 90".to_string(),
            parameters: HashMap::new(),
            options: QueryOptions::default(),
            context: QueryContext::default(),
        };

        let result = engine.execute(query).await.unwrap();
        assert!(result.is_success());

        if let QueryResultData::Arrow(batches) = &result.data {
            assert!(!batches.is_empty());
            let batch = &batches[0];
            assert_eq!(batch.num_rows(), 1); // 只有Alice的分数>90
            assert_eq!(batch.num_columns(), 2); // name和score两列
        } else {
            panic!("期望Arrow数据");
        }
    }

    #[tokio::test]
    async fn test_invalid_sql_query() {
        let engine = QueryEngine::new();

        let query = Query {
            id: "invalid_query".to_string(),
            query_type: QueryType::Sql,
            content: "SELECT * FROM non_existent_table".to_string(),
            parameters: HashMap::new(),
            options: QueryOptions::default(),
            context: QueryContext::default(),
        };

        let result = engine.execute(query).await.unwrap();
        assert!(!result.is_success());
        assert!(result.error.is_some());
    }

    #[tokio::test]
    async fn test_unsupported_query_type() {
        let engine = QueryEngine::new();

        let query = Query {
            id: "graphql_query".to_string(),
            query_type: QueryType::GraphQL,
            content: "{ users { id name } }".to_string(),
            parameters: HashMap::new(),
            options: QueryOptions::default(),
            context: QueryContext::default(),
        };

        let result = engine.execute(query).await.unwrap();
        assert!(!result.is_success());
        assert!(result.error.is_some());
        assert!(result.error.as_ref().unwrap().contains("不支持的查询类型"));
    }

    #[tokio::test]
    async fn test_get_table_schema() {
        let mut engine = QueryEngine::new();

        let data = vec![
            {
                let mut row = HashMap::new();
                row.insert("id".to_string(), json!(1));
                row.insert("name".to_string(), json!("test"));
                row.insert("score".to_string(), json!(95.5));
                row
            },
        ];

        engine.create_memory_table("test_table", data).await.unwrap();

        let schema = engine.get_table_schema("test_table").await.unwrap();
        assert_eq!(schema.fields().len(), 3);

        let field_names: Vec<_> = schema.fields().iter().map(|f| f.name().as_str()).collect();
        assert!(field_names.contains(&"id"));
        assert!(field_names.contains(&"name"));
        assert!(field_names.contains(&"score"));
    }
}
